package utils

import (
	"errors"
	"github.com/throttled/throttled"
	"github.com/throttled/throttled/store/memstore"
)

type RateLimiter struct {
	rateMap map[string]*throttled.GCRARateLimiter
}

// 项目需要初始化一个全局的rater
var Rater *RateLimiter

var limitPerSecMap = map[string]int{
	GetUserInfoRateLimiter: 2,
}

func init() {
	var errRater error
	// TODO 如果每种接口的限流规则不一样，里面的参数可以用一个配置替代
	Rater, errRater = NewRateLimiter(limitPerSecMap[GetUserInfoRateLimiter])
	if errRater != nil {
		// TODO 初始化错误直接panic,看个人业务需求
		panic(errRater)
	}
}

func NewRateLimiter(limitPerSec int) (*RateLimiter, error) {
	// 1、store
	store, err := memstore.New(65536)
	if err != nil {
		return nil, errors.New(err.Error())
	}
	// 2、初始化 quota
	// 为请求get_user_info专门设置一个
	rateUserInfo := throttled.RateQuota{
		MaxRate:  throttled.PerSec(limitPerSec), // 每秒的限额
		MaxBurst: 0,
	}
	// 3、 new GCRA
	instanceUserInfo, errInUserInfo := throttled.NewGCRARateLimiter(store, rateUserInfo)
	if errInUserInfo != nil {
		return nil, errors.New(errInUserInfo.Error())
	}

	// TODO 有其他限流规则重复步骤2与3

	// 构建ret 根据不同的规则灵活设定
	ret := &RateLimiter{rateMap: make(map[string]*throttled.GCRARateLimiter, 1)}
	ret.rateMap[GetUserInfoRateLimiter] = instanceUserInfo

	return ret, nil
}

func CheckRateLimit(key string) (bool, error) {
	rater, ok := Rater.rateMap[key]
	if !ok {
		return false, errors.New("CheckRateLimit-KeyError")
	}
	// 被限制住返回true
	bl, _, err := rater.RateLimit(key, 1)
	if err != nil {
		return false, errors.New("CheckRateLimit-RateError")
	}
	// 出发了ratelimit不想让业务继续执行！
	// 这里取反返回,函数返回的true表示没有触发限流
	return !bl, nil
}
